src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py (2)

216-216: extensions_mapping is now unused.

The self.extensions_mapping dictionary is no longer used after this refactor. Previously it was used to map extensions to keys for imagehandler, but now _convert_to_tensor always passes "jpg" directly. Consider removing the unused attribute to avoid confusion.

🧹 Proposed cleanup

     def __init__(self, imagespec):
         self.extensions = ["jpgs", "mp4s", "videos"]
-        self.extensions_mapping = {"jpgs": "jpg", "mp4s": "jpg", "videos": "jpg"}
         self.image_handler = imagehandler(imagespec)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py` at line
216, Remove the unused attribute self.extensions_mapping from the TaskEncoder
class: locate where it is defined (self.extensions_mapping = {"jpgs": "jpg",
"mp4s": "jpg", "videos": "jpg"}) and delete that line since _convert_to_tensor
now always passes "jpg" directly and imagehandler no longer relies on it; ensure
no other code references extensions_mapping (search for extensions_mapping)
before committing.

---

219-231: Add type hints for the new method.

Per coding guidelines, functions should have type hints for arguments and return types. This improves readability and enables static type checking.

💡 Suggested type hints

-    def _convert_to_tensor(self, data):
-        """Convert numpy array or bytes to tensor.
-        
-        The wds conversion script stores images as numpy arrays (HWC, uint8),
-        so we need to handle both numpy arrays and raw bytes.
-        """
+    def _convert_to_tensor(self, data: np.ndarray | bytes) -> torch.Tensor:
+        """Convert numpy array or bytes to tensor.
+
+        The wds conversion script stores images as numpy arrays (HWC, uint8),
+        so we need to handle both numpy arrays and raw bytes.
+
+        Args:
+            data: Image data as numpy array (HWC, uint8) or raw bytes.
+
+        Returns:
+            Tensor in CHW format with float32 values in [0, 1].
+        """
         if isinstance(data, np.ndarray):

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py` around
lines 219 - 231, Add type hints to _convert_to_tensor: annotate the data
parameter as Union[np.ndarray, bytes] and the return type as torch.Tensor (i.e.,
def _convert_to_tensor(self, data: Union[np.ndarray, bytes]) -> torch.Tensor).
Also ensure typing Union is imported (from typing import Union) and that
image_handler's return type is compatible with torch.Tensor; update
image_handler signature or cast its result to torch.Tensor if needed. This
change applies to the _convert_to_tensor method in task_encoder.py and any
related image_handler definition used by this method.